import os
import json
from functools import partial

import numpy as np

from tools.dataset_helper import open_img_np, plot_cmd
from tools.select_patch import crn_cmd


def open_QXSLAB(
    src_file, *, scale=0, mean=False, dims=True, **kwargs
):
    """
    open QXSLAB image as numpy ndarray

    Param
    -----
    scale : change to range of [0 scale] after normalize
            no change range if scale == 0
    mean :  return mean of channels
    dims :  keep dims == 3

    Note
    -----
    plt.imread read png as float with [0, 1]
    otherwise int 
    """
    if scale == 0:
        max_ = np.array([0., 0., 0.])
    else:
        max_ = np.array([1., 1., 1.])

    img_np = open_img_np(
        src_file, max_=max_, scale=scale, mean=mean, dims=dims, **kwargs
    )

    return img_np


def get_QXSLAB_filelist(dset_folder="", json_folder=""):
    '''
    将QXSLAB的全部图片分场景写入json文件
    '''

    os.makedirs(json_folder, exist_ok=True)

    city_range = ["QingDao", 3784, "ShangHai", 11949, "SanDiego", 20000]
    opt_folder = "opt_256_oc_0.2"
    sar_folder = "sar_256_oc_0.2"

    cnt = 1
    for i in range(len(city_range) // 2):
        city = city_range[2*i]
        ranges = city_range[2*i+1]
        cur_dict = {
            "OPT" : [], 
            "SAR" : []
        }

        while cnt <= ranges:
            cur_dict["OPT"].append(
                os.path.join(opt_folder, str(cnt)+".png")
            )
            cur_dict["SAR"].append(
                os.path.join(sar_folder, str(cnt)+".png")
            )

            cnt += 1

        with open(os.path.join(json_folder, city+".json"), "w") as fp:
            json.dump(cur_dict, fp, indent="\t")


def fn2wkt_QXSLAB(path: str) -> str:
    sample = os.path.basename(path)
    sample = os.path.splitext(sample)[0]
    wkt = sample
    
    return wkt


def wkt2fn_QXSLAB(wkt: str, modal:str) -> str:
    folder_dict = {
        "opt" : "opt_256_oc_0.2", 
        "sar" : "sar_256_oc_0.2"
    }

    fn = os.path.join(folder_dict[modal], wkt+".png")

    return fn


# write all files into list json per city
if __name__ == "__main__1":

    dset_folder = "E:\datasets\QXSLAB_SAROPT"
    json_folder = "E:\workspace\SOMatch\json/QXSLAB_SAROPT/"
    get_QXSLAB_filelist(dset_folder=dset_folder, json_folder=json_folder)


# test open_QXSLAB
if __name__ == "__main__1":

    # png = "E:/datasets/QXSLAB_SAROPT/sar_256_oc_0.2/4.png"
    png = "E:/datasets/QXSLAB_SAROPT/opt_256_oc_0.2/1.png"
    img = open_QXSLAB(png, scale=0)
    print(img.shape, img.max())


# work as cmd 
# "crn" : get crn 
if __name__ == "__main__":

    p = input("select program --> ")

    if p == "crn":

        crn_cmd(
            read_func=partial(open_QXSLAB, scale=255, mean=True), 
            fn2wkt=fn2wkt_QXSLAB, wkt2fn=wkt2fn_QXSLAB, 
            pt_range=[26, 230, 26, 230], 
        )

        # python -m tools.QXSLAB_helper -d "E:\datasets\QXSLAB_SAROPT" -s "E:\workspace\SOMatch\json\QXSLAB_QingDao_harris" -l "E:\workspace\SOMatch\json\QXSLAB_SAROPT\QingDao.json"   

    if p == "plot":

        plot_cmd(
            wkt2fn=wkt2fn_QXSLAB, 
            read_func=partial(open_QXSLAB, scale=0, mean=True)
        )

        # python -m tools.QXSLAB_helper -d "E:\datasets\QXSLAB_SAROPT" -o "E:\workspace\SOMatch\json\QXSLAB_QingDao_harris\OPT_HARRIS0.05.json" -s "E:\workspace\SOMatch\json\QXSLAB_QingDao_harris\SAR_HARRIS0.05.json" -v "E:\datasets\client-data\QXSLAB_QingDao_plot"

